#用来批量修正 YOLOv8-pose 标签目录，使每条标签包含相同数量的关键点

import os
import argparse

# —— 在这里直接修改默认的标签目录路径 —— 
DEFAULT_LABELS_DIR = r"C:\Users\lenovo\Downloads\rune_blender_opensource\render_video\labels"

def pad_keypoints_in_dir(labels_dir: str, kpt_num: int):
    """
    将 labels_dir 中所有 .txt 文件的关键点数补齐到 kpt_num，
    用 0 0 0 填充不足的部分，原地覆盖写回。
    """
    if not os.path.isdir(labels_dir):
        print(f"[ERROR] `{labels_dir}` 不是有效目录，已跳过。")
        return

    txt_files = [f for f in os.listdir(labels_dir) if f.lower().endswith('.txt')]
    if not txt_files:
        print(f"[WARN] 目录 `{labels_dir}` 下未找到 .txt 文件，已跳过。")
        return

    print(f"[INFO] 开始处理目录 `{labels_dir}` 中的 {len(txt_files)} 个标签文件…")
    for fname in txt_files:
        path = os.path.join(labels_dir, fname)
        lines_out = []
        with open(path, 'r', encoding='utf-8') as f:
            for line in f:
                parts = line.strip().split()
                if len(parts) < 5:
                    continue  # 非法或空行
                cls_bbox = parts[:5]
                kp = parts[5:]
                orig_kp_count = len(kp) // 3
                pad_count = max(0, kpt_num - orig_kp_count)
                kp_padded = kp + ['0','0','0'] * pad_count
                kp_padded = kp_padded[:kpt_num * 3]
                lines_out.append(' '.join(cls_bbox + kp_padded))

        with open(path, 'w', encoding='utf-8') as f:
            f.write('\n'.join(lines_out))
        print(f"  ✓ `{fname}` → 已补齐至 {kpt_num} 个关键点")

    print(f"[INFO] 目录 `{labels_dir}` 处理完毕。\n")

if __name__ == '__main__':
    parser = argparse.ArgumentParser(
        description='批量修正 YOLOv8-pose 标签目录，使每条标签包含相同数量的关键点'
    )
    parser.add_argument(
        '--dir', '-d',
        type=str,
        default=DEFAULT_LABELS_DIR,
        help=f'要处理的标签目录路径（默认: {DEFAULT_LABELS_DIR}）'
    )
    parser.add_argument(
        '--kpt-num', '-k',
        type=int,
        default=8,
        help='每条标注应包含的关键点数量'
    )
    args = parser.parse_args()

    print(f"配置：\n  要处理的目录：{args.dir}\n  目标关键点数：{args.kpt_num}\n")
    pad_keypoints_in_dir(args.dir, args.kpt_num)

    print("处理完毕。")
